# TODO: 进行STDN网络的构建
# DATE: 2023/3/17
# AUTHOR: Cheng Ze WUST

import torch
import torch.nn as nn
from nets.densenet import densenet169


class STDN(nn.Module):
    def __init__(self, num_classes, pretrained=False, net_name='densenet'):
        super(STDN, self).__init__()
        self.num_classes = num_classes
        if net_name == 'densenet':
            self.backbone = densenet169(pretrained)

        # 每个特征层的每个特征点对应的先验框数量为8
        mbox = 8
        # 经过STM模块输出的6个特征图尺寸
        outchannel = [800, 960, 1120, 1280, 360, 104]

        self.pool1 = nn.AvgPool2d(kernel_size=9)
        self.pool2 = nn.AvgPool2d(kernel_size=3)
        self.pool3 = nn.AvgPool2d(kernel_size=2, padding=1)

        # scale-transfer layer 超分上采样方法
        # 2x 1440,9,9 -> 360,18,18
        # 4x 1664,9,9 -> 104,36,36
        self.ps1 = nn.PixelShuffle(2)
        self.ps2 = nn.PixelShuffle(4)

        loc_layers = []
        conf_layers = []
        for i in range(6):
            loc_layers += [nn.Conv2d(outchannel[i], mbox * 4, kernel_size=3, padding=1)]
            conf_layers += [nn.Conv2d(outchannel[i], mbox * num_classes, kernel_size=3, padding=1)]

        self.loc = nn.ModuleList(loc_layers)
        self.conf = nn.ModuleList(conf_layers)

    def forward(self, x):
        sources = list()
        loc = list()
        conf = list()
        # 获取经过backbone前向传播获得的6个特征图
        _, feature = self.backbone(x)

        out1 = self.pool1(feature[0])   # 9×9 平均池化
        out2 = self.pool2(feature[1])   # 3×3 平均池化
        out3 = self.pool3(feature[2])   # 2×2 平均池化
        out4 = feature[3]               # 不作处理
        out5 = self.ps1(feature[4])     # 2× 超分上采样
        out6 = self.ps2(feature[5])     # 4× 超分上采样

        sources.append(out1)
        sources.append(out2)
        sources.append(out3)
        sources.append(out4)
        sources.append(out5)
        sources.append(out6)

        # 为获得的6个有效特征层添加回归预测和分类预测
        for (x, l, c) in zip(sources, self.loc, self.conf):
            loc.append(l(x).permute(0, 2, 3, 1).contiguous())
            conf.append(c(x).permute(0, 2, 3, 1).contiguous())

        # 进行reshape方便堆叠
        # loc会reshape到batch_size, num_anchors, 4
        # conf会reshap到batch_size, num_anchors, self.num_classes
        loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
        conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)

        output = (
            loc.view(loc.size(0), -1, 4),
            conf.view(conf.size(0), -1, self.num_classes),
        )

        return output


if __name__ == "__main__":
    model = STDN(num_classes=1000)
    # 查看网络结构
    print(model)
    # 查看网络各层参数
    from torchstat import stat
    # stat(model, input_size=(3, 300, 300))
